import argparse
from pathlib import Path
import re
from typing import ClassVar
import polars as pl

from pre_base import DATA_DIR, PrepareBase


class PrepareClingen(PrepareBase):
    name: ClassVar[str] = "ClinGen"
    header: ClassVar[str] = '##INFO=<ID=CLINGEN,Number=1,Type=String,Description="ClinGen, chrom:start:end:hi_score:ts_score">'

    def read_input(self):
        dataframe = pl.read_csv(
            self.input_file,
            separator="\t",
            skip_rows=5,
            columns=["#ISCA ID", "Genomic Location", "Haploinsufficiency Score", "Triplosensitivity Score"],
            new_columns=["ISCA", "Location", "HI", "TS"],
        )
        dataframe = (
            dataframe.with_columns(
                pl.col("Location").map_elements(lambda x: dict(zip(["Chromosome", "Start", "End"], re.split("[:-]", x)[0:3])), pl.Struct)
            )
            .unnest(["Location"])
            .with_columns(pl.col("Start").cast(pl.Int64).alias("Start"), pl.col("End").cast(pl.Int64).alias("End"))
            .with_columns(
                pl.format(
                    "{}:{}:{}:{}:{}",
                    pl.col("Chromosome"),
                    pl.col("Start"),
                    pl.col("End"),
                    pl.col("HI"),
                    pl.col("TS"),
                ).alias("INFO")
            )
            .filter(pl.col("INFO").is_not_null())
            .sort(["Chromosome", "Start", "End"])
            .select(["Chromosome", "Start", "End", "INFO"])
        )
        return dataframe


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare ClinGen")
    parser.add_argument("--input_file", "-i", default=DATA_DIR / "raw" / "ClinGen_region_curation_list_GRCh38.tsv", type=Path, help="input file")
    parser.add_argument("--outdir", "-o", default=DATA_DIR / "bed", type=Path, help="output directory")
    parser.set_defaults(func=lambda args: PrepareClingen.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)

    # f"{variant_id}:{self.record.chrom}:{start}:{end}:{revstat}:{clnsig}"
